﻿#include "PassFlow.h"
#include "Application.h"
#include "EulerEquation.h"
#include "Mesh.h"
#include "PassFlow.h"
#include "Timer.h"
#include "Transient.h"
#include <vtkm/Bounds.h>
#include <vtkm/Geometry.h>
#include <vtkm/cont/CellLocatorGeneral.h>
#include <vtkm/cont/CellLocatorTwoLevel.h>

struct PassFlowIC : WorkletInitCondition<PassFlowIC>
{
  VTKM_EXEC ValueType ComputeValue(const Vec3f& p) const
  {
    ValueType primary_value;
    primary_value = ValueType{ 1.4, 3, 0, 0, 1 };

    return PrimaryToConservation(primary_value);
  }
  VTKM_EXEC ValueType PrimaryToConservation(const ValueType& primary_value) const
  {
    FloatDefault gam = 1.4, gamm1 = gam - 1;
    ValueType value;
    auto r = primary_value[0];
    auto u = primary_value[1];
    auto v = primary_value[2];
    auto w = primary_value[3];
    auto p = primary_value[4];

    value[0] = r;
    value[1] = r * u;
    value[2] = r * v;
    value[3] = r * w;
    value[4] = p / gamm1 + 0.5 * r * (u * u + v * v);
    return value;
  }
};

struct ComputeResidual_PassFlow : vtkm::worklet::WorkletCellNeighborhood
{
  VTKM_CONT
  ComputeResidual_PassFlow(FloatDefault t)
    : _t(t)
  {
  }

  using ControlSignature = void(CellSetIn,
                                FieldInNeighborhood solution,
                                FieldOut Rhs,
                                FieldIn centroid,
                                FieldIn spacing,
                                FieldIn blk_id,
                                FieldIn bnd_id,
                                ExecObject locator,
                                WholeArrayIn sln);
  using ExecutionSignature = void(Boundary, _2, _3, _4, _5, _6, _7, _8, _9);

  template<typename NeighIn, typename RhsType, typename LocatorType, typename FieldPortal>
  VTKM_EXEC void operator()(const vtkm::exec::BoundaryState& boundary,
                            const NeighIn& value,
                            RhsType& rhs,
                            const Vec3f& centroid,
                            const Vec3f& spacing,
                            const vtkm::Int16& blk_id,
                            const Vec<vtkm::Int16, 4>& bnd_id,
                            const LocatorType& locator,
                            const FieldPortal& sln) const
  {
    if (blk_id == -1)
    {
      rhs[0] = 0;
      rhs[1] = 0;
      rhs[2] = 0;
      rhs[3] = 0;
      rhs[4] = 0;
      return;
    }
    Vec3f normal{ 0, 0, 0 };
    DataPack<RhsType::NUM_COMPONENTS> dl, dr;
    dl.uh = value.Get(0, 0, 0);
    dl.InvisTerm();
    {
      normal = Vec3f{ -1, 0, 0 };
      dr.uh = value.Get(-1, 0, 0);
      RhsType lflux, rflux;
      Vec3f p = centroid;
      p[0] = p[0] - 0.5 * spacing[0];
      BoundaryCondition(bnd_id[3], _t, p, normal, dl, dr, locator, sln);
      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, lflux);

      dr.uh = value.Get(1, 0, 0);
      normal = Vec3f{ 1, 0, 0 };
      p = centroid;
      p[0] = p[0] + 0.5 * spacing[0];
      BoundaryCondition(bnd_id[1], _t, p, normal, dl, dr, locator, sln);
      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, rflux);

      rhs = (rflux + lflux) / spacing[0];
    }

    {
      dr.uh = value.Get(0, -1, 0);
      RhsType lflux, rflux;
      normal = Vec3f{ 0, -1, 0 };
      Vec3f p = centroid;
      p[1] = p[1] - 0.5 * spacing[1];
      BoundaryCondition(bnd_id[0], _t, p, normal, dl, dr, locator, sln);
      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, lflux);

      dr.uh = value.Get(0, 1, 0);
      normal = Vec3f{ 0, 1, 0 };
      p = centroid;
      p[1] = p[1] + 0.5 * spacing[1];
      BoundaryCondition(bnd_id[2], _t, p, normal, dl, dr, locator, sln);
      dr.InvisTerm();
      RiemannFlux(dl, dr, normal, rflux);

      rhs += (rflux + lflux) / spacing[1];
    }
  }

  template<typename DataPackType, typename VType>
  VTKM_EXEC void RiemannFlux(const DataPackType& dl,
                             const DataPackType& dr,
                             const Vec3f& normal,
                             VType& riemann_flux) const
  {
    for (IdComponent eq = 0; eq < VType::NUM_COMPONENTS; eq++)
    {
      riemann_flux[eq] =
        vtkm::Dot(dl.inv_term[eq] + dr.inv_term[eq], normal) / 2 + 1.5 * (dl.uh[eq] - dr.uh[eq]);
    }
  }

  template<typename DataPackType, typename LocatorType, typename FieldPortal>
  VTKM_EXEC void BoundaryCondition(const vtkm::Int16& bid,
                                   const FloatDefault& t,
                                   const Vec3f& p,
                                   const Vec3f& normal,
                                   const DataPackType& dl,
                                   DataPackType& dr,
                                   const LocatorType& locator,
                                   const FieldPortal& sln) const
  {
    if (bid == 1) // 出口
    {
      auto& ul = dl.uh;
      auto& ur = dr.uh;
      ur[0] = ul[0];
      ur[1] = ul[1];
      ur[2] = ul[2];
      ur[3] = ul[3];
      ur[4] = ul[4];
    }
    if (bid == 3) // 进口
    {
      using ValueType = Vec<FloatDefault, 5>;
      ValueType primary_value;
      primary_value = ValueType{ 1.4, 3, 0, 0, 1 };

      FloatDefault gam = 1.4, gamm1 = gam - 1;
      ValueType value;
      auto r = primary_value[0];
      auto u = primary_value[1];
      auto v = primary_value[2];
      auto w = primary_value[3];
      auto pre = primary_value[4];

      auto& ur = dr.uh;

      ur[0] = r;
      ur[1] = r * u;
      ur[2] = r * v;
      ur[3] = r * w;
      ur[4] = pre / gamm1 + 0.5 * r * (u * u + v * v);
    }
    if (bid == 0 || bid == 2 || bid == 4) //wall
    {
      auto& ul = dl.uh;
      auto& ur = dr.uh;
      Vec3f mom{ ul[1], ul[2], ul[3] };
      Real vn = vtkm::Dot(mom, normal);
      ur[0] = ul[0];
      ur[1] = ul[1] - 2 * vn * normal[0];
      ur[2] = ul[2] - 2 * vn * normal[1];
      ur[3] = ul[3] - 2 * vn * normal[2];
      ur[4] = ul[4];
    }
    if (bid == 11)
    {
      auto& ul = dl.uh;
      auto& ur = dr.uh;
      Id cell_id;
      Vec3f pcoords;
      vtkm::ErrorCode status = locator.FindCell(p, cell_id, pcoords);
      ur = sln.Get(cell_id);
      //std::cout << cell_id << std::endl;
    }
  }

  FloatDefault _t;
};

PassFlow::PassFlow(const Configuration& cfg)
  : EulerEquation(cfg)
{
}

void PassFlow::Init()
{
  EulerEquation::Init();
  //vtkm::cont::CellLocatorTwoLevel locator;
  //locator.SetDensityL1(100);
  //locator.SetDensityL2(10);
}

void PassFlow::InitialCondition()
{
  vtkm::cont::Timer timer;
  timer.Start();
  vtkm::cont::Invoker invoker;
  invoker(PassFlowIC{}, _mesh.GetCellSet(), _solution, _mesh.GetCoordinate());


  Success("初始条件设置完成", timer.GetElapsedTime(), "s");
}

void PassFlow::UpdateResidual()
{
  auto& transient = static_cast<Transient&>(*(_app.GetExecutioner()));
  Real current_time = transient.CurrentTime();

  Timer timer;
  timer.Start();
  const auto& centroid = _mesh.GetCellField("center");
  const auto& spacing = _mesh.GetCellField("spacing");
  auto blk_id = _mesh.GetCellField("blk_id").GetData().AsArrayHandle<ArrayHandle<vtkm::Int16>>();
  auto bnd_id =
    _mesh.GetCellField("bnd_id").GetData().AsArrayHandle<ArrayHandle<Vec<vtkm::Int16, 4>>>();


  vtkm::cont::CellLocatorGeneral _locator;
  auto& cell_set = _other->GetDateSet().GetCellSet();
  auto& coord = _other->GetDateSet().GetCoordinateSystem();
  _locator.SetCellSet(cell_set);
  _locator.SetCoordinates(coord);
  _locator.Update();


  vtkm::cont::Invoker{}(ComputeResidual_PassFlow{ current_time },
                        _mesh.GetCellSet(),
                        _solution,
                        _rhs,
                        centroid,
                        spacing,
                        blk_id,
                        bnd_id,
                        _locator,
                        _other->GetSolution());



  Info("ForwardStep::UpdateResidual运行完成，", "时间: ", timer.GetElapsedTime(), " s");
}

void PassFlow::SetOther(std::shared_ptr<PassFlow> other)
{
  _other = other;
}
